# Kestrel: An XMPP-based Many-Task Computing Scheduler
# Author: Lance Stout <lancestout@gmail.com>
#
# Copyright 2009-2010 Clemson University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Plotter module.

The plotter module allows for displaying the number of agents online
in the system over time and their status, or the number of completed
tasks over time.
"""

import yaml


class Plotter(object):
    """
    The plotter module allows for displaying the number of agents online
    in the system over time and their status, or the number of completed
    tasks over time.

    Events listened for by this module:
      log/profile/received
      log/presence/offline
      log/presence/available
      analyzer/read/complete
    """

    def __init__(self, kernel):
        """
        Instantiate a new plot.
        """
        self.kernel = kernel
        self.events = {'log/profile/received': self.profile_received,
                       'log/presence/offline': self.offline,
                       'log/presence/available': self.available,
                       'log/presence/busy': self.busy,
                       'log/job/approved': self.job_start,
                       'log/job/finished': self.job_finish,
                       'log/presence/send': self.presence_send,
                       'log/presence/subscribe': self.presence_recv,
                       'log/presence/subscribed': self.presence_recv,
                       'log/message/send': self.message_send,
                       'log/message': self.message_recv,
                       'log/job/task/start': self.task_start,
                       'log/job/task/end': self.task_end,
                       'analyzer/read/complete': self.finish}
        self.kernel.register(self)
        self.profiles = {}
        self.plot = {}
        self.agents = {}
        self.running = {'available': 0,
                        'busy': 0,
                        'profiles': 0,
                        'offline': 0,
                        'jobs': 0,
                        'tasks_running': 0,
                        'tasks_complete': 0,
                        'presence_sent': 0,
                        'presence_recv': 0,
                        'messages_sent': 0,
                        'messages_recv': 0}

    def prep_time(self, time):
        """
        Given a time instant, set the number of profiles and available
        agents to the current running totals.
        """
        if self.plot.get(time, None) is None:
            self.plot[time] = {
                'profiles': self.running['profiles'],
                'available': self.running['available'],
                'busy': self.running['busy'],
                'offline': 0,
                'jobs': self.running['jobs'],
                'tasks_running': self.running['tasks_running'],
                'tasks_complete': self.running['tasks_complete'],
                'presence_sent': self.running['presence_sent'],
                'presence_recv': self.running['presence_recv'],
                'messages_sent': self.running['messages_sent'],
                'messages_recv': self.running['messages_recv']}

    def strip_time(self, time):
        """Remove the dot from a timestamp."""
        t1 = time[0:time.index('.')]
        t2 = time[time.index('.') + 1:]
        return t1 + t2

    def profile_received(self, data=None):
        """
        Increase the number of worker profiles received.
        """
        time = data['time']
        data = data['data']
        data = yaml.load(data)

        self.prep_time(time)

        agent = data['source']
        if agent not in self.agents:
            self.agents[agent] = {'profile': False, 'status': 'offline'}

        if not self.agents[agent]['profile']:
            self.running['profiles'] += 1
            self.plot[time]['profiles'] = self.running['profiles']
            self.profiles[data['source']] = False
            self.agents[agent]['profile'] = True

    def available(self, data=None):
        """
        Increase the number of available agents.

        If multiple, consecutibe available presences are received for
        a single JID, only one is counted.
        """
        time = data['time']
        data = data['data']
        data = yaml.load(data)

        self.prep_time(time)

        agent = data['full']
        if agent not in self.agents:
            self.agents[agent] = {'profile': False, 'status': 'offline'}

        prev_status = self.agents[agent]['status']
        if prev_status != 'available':
            self.running[prev_status] -= 1
            self.running['available'] += 1
            self.plot[time]['available'] = self.running['available']
            self.plot[time][prev_status] = self.running[prev_status]
            self.agents[agent]['status'] = 'available'

        self.presence_recv(data)

    def busy(self, data=None):
        """
        Increase the number of busy agents.
        """
        time = data['time']
        data = data['data']
        data = yaml.load(data)

        self.prep_time(time)

        agent = data['full']
        if agent not in self.agents:
            self.agents[agent] = {'profile': False, 'status': 'offline'}

        prev_status = self.agents[agent]['status']
        if prev_status != 'busy':
            self.running[prev_status] -= 1
            self.running['busy'] += 1
            self.plot[time]['busy'] = self.running['busy']
            self.plot[time][prev_status] = self.running[prev_status]
            self.agents[agent]['status'] = 'busy'

        self.presence_recv(data)

    def offline(self, data=None):
        """
        Decrease the number of online agents.

        If multiple, consecutibe offline presences are received for
        a single JID, only one is counted.
        """
        time = data['time']
        data = data['data']
        data = yaml.load(data)
        print data

        self.prep_time(time)

        agent = data['full']
        if agent not in self.agents:
            self.agents[agent] = {'profile': False, 'status': 'offline'}

        prev_status = self.agents[agent]['status']
        if prev_status != 'offline':
            self.running[prev_status] -= 1
            self.running['offline'] += 1
            self.plot[time]['offline'] = self.running['offline']
            self.plot[time][prev_status] = self.running[prev_status]
            self.agents[agent]['status'] = 'offline'

        self.presence_recv(data)

    def job_start(self, data=None):
        """
        Mark a job as started.
        """
        time = data['time']
        self.prep_time(time)
        self.running['jobs'] += 1
        self.plot[time]['jobs'] = self.running['jobs']

    def job_finish(self, data=None):
        """
        Mark a job as completed.
        """
        time = data['time']
        self.prep_time(time)
        self.running['jobs'] -= 1
        self.plot[time]['jobs'] = self.running['jobs']

    def presence_recv(self, data=None):
        """
        Increase number of presence notifications received.
        """
        time = data['time']
        self.prep_time(time)
        self.running['presence_recv'] += 1
        self.plot[time]['presence_recv'] = self.running['presence_recv']

    def task_start(self, data=None):
        """
        Increase number of tasks running.
        """
        time = data['time']
        self.prep_time(time)
        self.running['tasks_running'] += 1
        self.plot[time]['tasks_running'] = self.running['tasks_running']

    def task_end(self, data=None):
        """
        Decrease number of tasks running, increase number of tasks completed.
        """
        time = data['time']
        self.prep_time(time)
        self.running['tasks_running'] -= 1
        self.plot[time]['tasks_running'] = self.running['tasks_running']
        self.running['tasks_complete'] += 1
        self.plot[time]['tasks_complete'] = self.running['tasks_complete']

    def presence_send(self, data=None):
        """
        Increase number of presence notifications sent.
        """
        time = data['time']
        self.prep_time(time)
        self.running['presence_sent'] += 1
        self.plot[time]['presence_sent'] = self.running['presence_sent']

    def message_send(self, data=None):
        """
        Increase number of messages sent.
        """
        time = data['time']
        self.prep_time(time)
        self.running['messages_sent'] += 1
        self.plot[time]['messages_sent'] = self.running['messages_sent']

    def message_recv(self, data=None):
        """
        Increase number of messages received.
        """
        time = data['time']
        self.prep_time(time)
        self.running['messages_recv'] += 1
        self.plot[time]['messages_recv'] = self.running['messages_recv']

    def finish(self, data=None):
        """
        Display a plot showing the data collected from the log file.
        """
        from pylab import plotfile, show

        f = open('plot.tmp', 'w+')
        times = self.plot.keys()
        times.sort()

        start = int(self.strip_time(times[0]))
        prev = (start, 0, 0)

        for time in times:
            # Convert time from hundredths of seconds to seconds
            t = (int(self.strip_time(time)) - start) / 100
            p = self.plot[time]['profiles']
            a = self.plot[time]['available']
            b = self.plot[time]['busy']
            j = self.plot[time]['jobs']
            pr = self.plot[time]['presence_recv']
            ps = self.plot[time]['presence_sent']
            mr = self.plot[time]['messages_recv']
            ms = self.plot[time]['messages_sent']
            tr = self.plot[time]['tasks_running']
            tc = self.plot[time]['tasks_complete']

            f.write('%0.2f, %d, %d, %d, %d, %d, %d, %d, %d, %d, %d\n' %
                    (t, p, a, b, j, pr, ps, mr, ms, tr, tc))
            prev = (t, p, a, b, j, pr, ps, mr, ms, tr, tc)

        # Write an additional entry with 0 available agents
        # because the graph doesn't want to show the last data point
        f.write('%0.2f, %d, %d, %d, %d, %d, %d, %d, %d, %d, %d\n' %
                (prev[0] + 1,
                 prev[1],
                 prev[2],
                 prev[3],
                 prev[4],
                 prev[5],
                 prev[6],
                 prev[7],
                 prev[8],
                 prev[9],
                 prev[10]))
        f.close()

        plotfile('plot.tmp',
                 (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10),
                 subplots=False,
                 names=['Time (Seconds)',
                        'Worker Profiles',
                        'Agents Available',
                        'Agents Busy',
                        'Queued Jobs',
                        'Presence Notifications (Recv)',
                        'Presence Notifications (Sent)',
                        'Messages (Recv)',
                        'Messages (Sent)',
                        'Tasks (Running)',
                        'Tasks (Complete)'])
        show()
